from digi_sign import *


def test_encryption(e, d, n):
    test_file = open('./test.message', 'rb')
    encrypted_bytes = rsa_encrypt(test_file.read(), e, n)
    test_file.close()

    with open('./test.message.encrypted', 'wb') as out:
        out.write(encrypted_bytes)

    with open('./test.message.encrypted', 'rb') as infile, open('./test.message.decrypted', 'wb') as outfile:
        outfile.write(rsa_decrypt(infile.read(), d, n))


def test_digi_sign(e, d, n):
    with open('./test.message', 'rb') as in_file:
        print("Reading input message.")
        msg = in_file.read()
        print("Input message has been read.")
    print("Signing input message.")
    signature = sign_message(msg, d, n)
    print("Input message signed.")
    print("Signature valid? ", is_signature_valid(msg, signature, e, n))


def main():
    e, n = rsa_read_public_key()
    d, _ = rsa_read_private_key()

    test_encryption(e, d, n)
    test_digi_sign(e, d, n)

    # print(rsa_encode(123, e, n))
    # print(rsa_encode(rsa_encode(123, e, n), d, n))


if __name__ == '__main__':
    main()
